package traceroute

import (
	"fmt"
	"net"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"time"
)

// TracerouteParam 路由追踪参数
// Host: 目标主机
// Port: 目标端口（0表示不指定端口，默认使用ICMP）
// MaxHop: 最大跳数
// Timeout: 超时时间（秒）
type TracerouteParam struct {
	Host    string
	Port    int
	MaxHop  int
	Timeout int
}

// TracerouteHop 路由追踪每一跳的信息
// Hop: 跳数
// Addr: 跳点地址
// RTT: 往返时间
// Success: 是否成功
// Err: 错误信息
// Hostname: 跳点主机名
// Port: 端口（如有）
type TracerouteHop struct {
	Hop      int
	Addr     string
	Hostname string
	RTT      time.Duration
	Success  bool
	Err      error
	Port     int
}

// TracerouteResult 路由追踪结果
// Hops: 跳点信息
// Success: 是否追踪到目标
// Target: 目标主机
// TargetPort: 目标端口
// TotalTime: 总耗时
// Error: 错误信息
type TracerouteResult struct {
	Hops       []TracerouteHop
	Success    bool
	Target     string
	TargetPort int
	TotalTime  time.Duration
	Error      error
}

// NewTracerouteParam 创建并初始化一个 TracerouteParam 实例
//
// 参数：
//
//	host: 目标主机
//	port: 目标端口（0表示不指定端口，默认ICMP）
//	maxHop: 最大跳数，<=0 时默认30
//	timeout: 超时时间（秒），<=0 时默认3秒
//
// 返回：
//
//	TracerouteParam 路由追踪参数实例
func NewTracerouteParam(host string, port, maxHop, timeout int) TracerouteParam {
	if maxHop <= 0 {
		maxHop = 30
	}
	if timeout <= 0 {
		timeout = 3
	}
	return TracerouteParam{
		Host:    host,
		Port:    port,
		MaxHop:  maxHop,
		Timeout: timeout,
	}
}

// Traceroute 路由追踪主函数，支持指定端口（端口为0时使用ICMP，否则使用UDP/TCP）
func (param TracerouteParam) Traceroute() TracerouteResult {
	var result TracerouteResult
	result.Target = param.Host
	result.TargetPort = param.Port
	startTime := time.Now()
	maxHop := param.MaxHop
	if maxHop <= 0 {
		maxHop = 30
	}
	timeout := param.Timeout
	if timeout <= 0 {
		timeout = 3
	}
	for ttl := 1; ttl <= maxHop; ttl++ {
		var hop TracerouteHop
		hop.Hop = ttl
		addr, rtt, err := probeHop(param.Host, param.Port, ttl, timeout)
		hop.RTT = rtt
		hop.Addr = addr
		hop.Port = param.Port
		hop.Success = (err == nil)
		hop.Err = err
		if addr != "" {
			host, _ := net.LookupAddr(addr)
			if len(host) > 0 {
				hop.Hostname = host[0]
			}
		}
		result.Hops = append(result.Hops, hop)
		if err == nil && addr == param.Host {
			result.Success = true
			break
		}
	}
	result.TotalTime = time.Since(startTime)
	return result
}

// TracerouteBySystem 调用系统traceroute/tracert命令，获取每一跳真实经过的IP节点
func (param TracerouteParam) TracerouteBySystem() TracerouteResult {
	var result TracerouteResult
	result.Target = param.Host
	result.TargetPort = param.Port
	startTime := time.Now()

	var cmd *exec.Cmd
	maxHop := param.MaxHop
	if maxHop <= 0 {
		maxHop = 30
	}
	timeout := param.Timeout
	if timeout <= 0 {
		timeout = 3
	}
	if runtime.GOOS == "windows" {
		cmd = exec.Command("tracert", "-d", "-h", fmt.Sprintf("%d", maxHop), "-w", fmt.Sprintf("%d", timeout*1000), param.Host)
	} else {
		args := []string{"-n", "-m", fmt.Sprintf("%d", maxHop), "-w", fmt.Sprintf("%d", timeout), param.Host}
		if param.Port > 0 {
			args = append(args, "-p", fmt.Sprintf("%d", param.Port))
		}
		cmd = exec.Command("traceroute", args...)
	}
	output, err := cmd.CombinedOutput()
	if err != nil {
		result.Error = err
		return result
	}
	lines := strings.Split(string(output), "\n")
	ipRegex := regexp.MustCompile(`(\d+\.\d+\.\d+\.\d+)`)
	targetReached := false
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		matches := ipRegex.FindAllString(line, -1)
		if len(matches) > 0 {
			ip := matches[0]
			hop := TracerouteHop{
				Hop:      len(result.Hops) + 1,
				Addr:     ip,
				Hostname: "",
				RTT:      0,
				Success:  true,
				Err:      nil,
				Port:     param.Port,
			}
			result.Hops = append(result.Hops, hop)
			if ip == param.Host {
				targetReached = true
				break
			}
		} else {
			// 没有IP，说明该跳失败，直接停止
			hop := TracerouteHop{
				Hop:      len(result.Hops) + 1,
				Addr:     "*",
				Hostname: "",
				RTT:      0,
				Success:  false,
				Err:      fmt.Errorf("timeout or unreachable"),
				Port:     param.Port,
			}
			result.Hops = append(result.Hops, hop)
			break
		}
	}
	result.TotalTime = time.Since(startTime)
	result.Success = targetReached
	return result
}

// probeHop 探测单跳，端口为0时用ICMP，否则用UDP/TCP
func probeHop(host string, port, ttl, timeout int) (string, time.Duration, error) {
	start := time.Now()
	var err error
	address := host
	if port > 0 {
		address = fmt.Sprintf("%s:%d", host, port)
	} else {
		address = fmt.Sprintf("%s:33434", host)
	}
	udpAddr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return "", 0, err
	}
	conn, err := net.DialUDP("udp", nil, udpAddr)
	if err != nil {
		return "", 0, err
	}
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	// 设置TTL
	if f, err := conn.SyscallConn(); err == nil {
		f.Control(func(fd uintptr) {
			_ = setUDPTTL(fd, ttl)
		})
	}
	// 发送探测包
	msg := []byte("traceroute_test")
	conn.Write(msg)

	// 尝试读取ICMP超时响应，获取每一跳的真实IP
	buf := make([]byte, 512)
	conn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	_, addr, err := conn.ReadFrom(buf)
	cost := time.Since(start)
	if err == nil {
		if udpAddr, ok := addr.(*net.UDPAddr); ok {
			return udpAddr.IP.String(), cost, nil
		}
		return addr.String(), cost, nil
	}
	return "", cost, err
}
